import datetime
import functools
import json
import logging
import os
import traceback
from concurrent.futures import ProcessPoolExecutor

import bokeh.embed
import google.protobuf.json_format as json_format
import pandas as pd
import pytz
from absl import app, flags

import coin.proto.coin_order_enums_pb2 as order_enums
from coin.base.datetime_util import convert_string_to_dates, iterate_date
from coin.proto.coin_strategy_pb2 import StrategyLog
from coin.proto.coin_order_gateway_pb2 import OrderGatewayLog, OrderEvent
from coin.proto.coin_request_pb2 import AccountRequestProto
from coin.support.accounting.logic.query_util import AccountRequest
from coin.support.pta.util.info_util import get_strat_info_list
from coin.support.proto_log.logic.util import run_from_strat_log_archive
from coin.support.pta.logic.pta_stats_calculator import get_order_sign
from coin.tool.strat_monitor.plot.strat_plotter import DailyStrategyPlotterConfig
from coin.tool.strat_monitor.util.monitor_util import get_strat_from_json_config
from coin.util.queue.config import KafkaConfig
from xunkemgmt_client.client.util.query_util import query_strategies


FLAGS = flags.FLAGS


class SingleOrderContext(object):
  def __init__(self):
    self.pending_time = None
    self.accepted_time = None
    self.dead_time = None
    self.order_state = order_enums.CREATED_ORDER


class SymbolOrderEventContainer(object):
  def __init__(self, account_request, symbol):
    self._account_request = account_request
    self._symbol = symbol
    self._working_orders = set()
    self._accm_num_error_orders = 0
    self._order_events = []
    self._order_contexts = {}

  @property
  def account_request(self):
    return self._account_request

  @property
  def symbol(self):
    return self._symbol

  def _get_single_order_context(self, internal_order_id):
    if internal_order_id not in self._order_contexts:
      self._order_contexts[internal_order_id] = SingleOrderContext()
    return self._order_contexts[internal_order_id]

  def update_by_order_event(self, ts, order_event):
    prev_num_working_orders = len(self._working_orders)
    order_dict = _gen_order_dict(order_event)
    internal_order_id = order_dict['internal_order_id']
    order_context = self._get_single_order_context(internal_order_id)
    if order_event.order_state == order_enums.PENDING_ORDER:
      assert order_context.pending_time is None
      assert order_context.accepted_time is None
      assert order_context.dead_time is None
      assert order_context.order_state == order_enums.CREATED_ORDER
      order_context.pending_time = order_event.event_time
      self._working_orders.add(internal_order_id)
    elif order_event.order_state == order_enums.WORKING_ORDER:
      assert order_context.order_state in (
          order_enums.CREATED_ORDER,
          order_enums.PENDING_ORDER,
          order_enums.WORKING_ORDER), order_context.order_state
      if order_context.order_state == order_enums.CREATED_ORDER:
        logging.debug(
            'working order not found in history. %s %s' % (ts, order_event))
      if order_context.order_state == order_enums.PENDING_ORDER:
        assert order_context.pending_time is not None
      if order_context.accepted_time is None:
        order_context.accepted_time = order_event.event_time
        if order_context.pending_time is not None:
          order_dict['order_pending_duration'] = \
              order_context.accepted_time - order_context.pending_time
      else:
        assert order_context.accepted_time <= order_event.event_time
      self._working_orders.add(internal_order_id)
    elif order_event.order_state == order_enums.DEAD_ORDER:
      if order_context.order_state == order_enums.CREATED_ORDER:
        logging.debug(
            'dead order not found in history. %s %s' % (ts, order_event))
      if order_context.dead_time is None:
        order_context.dead_time = order_event.event_time
      else:
        assert order_context.dead_time <= order_event.event_time
      if order_context.pending_time is not None:
        order_dict['order_life_duration'] = \
            order_context.dead_time - order_context.pending_time
      self._working_orders.discard(internal_order_id)
    else:
      raise ValueError('Unknown order state. %s' % order_event.order_state)
    order_context.order_state = order_event.order_state
    curr_num_working_orders = len(self._working_orders)
    if curr_num_working_orders != prev_num_working_orders:
      order_dict['num_working_orders'] = curr_num_working_orders
      order_dict['working_orders'] = ','.join(sorted(list(self._working_orders)))
    if order_event.type in (OrderEvent.ORDER_ERROR, OrderEvent.ORDER_REJECTED):
      self._accm_num_error_orders += 1
      order_dict['accm_num_error_orders'] = self._accm_num_error_orders
    if order_context.pending_time is not None:
      order_dict['pending_timestamp'] = order_context.pending_time
    self._order_events.append(order_dict)

  def get_history(self):
    return self._order_events


class OrderLogReplayer(object):
  def __init__(self, strat_info, start_date, end_date):
    self._strat_info = strat_info
    self._start_date = start_date
    self._end_date = end_date
    self._oe_containers = {}

  def on_log(self, ts, log):
    pb = StrategyLog()
    pb.ParseFromString(log)
    if pb.type == StrategyLog.OG_LOG:
      og_log = pb.og_log
      if og_log.type == OrderGatewayLog.ORDER_EVENT:
        assert og_log.HasField('account_request')
        order_event = og_log.event
        assert order_event.HasField('symbol')
        symbol_oe_container = self._get_symbol_oe_container(
            og_log.account_request, order_event.symbol)
        symbol_oe_container.update_by_order_event(ts, order_event)

  def _get_symbol_oe_container(self, account_request, symbol):
    assert isinstance(account_request, AccountRequestProto)
    acct_req = AccountRequest(market_type=account_request.market_type,
                              exchange=account_request.exchange,
                              owner=account_request.owner)
    key = (acct_req, symbol)
    if key not in self._oe_containers:
      self._oe_containers[key] = SymbolOrderEventContainer(account_request, symbol)
    return self._oe_containers[key]

  def dump_history(self, root_dir):
    for container in self._oe_containers.values():
      output_dir = get_order_event_data_dir(
          root_dir=root_dir,
          start_date=self._start_date,
          end_date=self._end_date,
          strategy_name=self._strat_info.strategy_name,
          machine=self._strat_info.machine,
          account_request=container.account_request,
          symbol=container.symbol)
      os.makedirs(output_dir, exist_ok=True)
      filename = 'order_event.json'
      with open('%s/%s' % (output_dir, filename), mode='w') as fd:
        json.dump(container.get_history(), fd)

  def dump_html(self, root_dir, server_url):
    html = """
<!DOCTYPE html>
<html>
<body>
%s
</body>
</html>
"""
    config = DailyStrategyPlotterConfig(self._start_date, self._end_date, root_dir)
    for container in self._oe_containers.values():
      strategy_name = self._strat_info.strategy_name
      machine=self._strat_info.machine
      output_dir = config.get_order_plot_dir(
          strategy_name=strategy_name,
          machine=machine,
          account_request=container.account_request,
          symbol=container.symbol)
      arguments = {
          'start_date': self._start_date.strftime('%Y-%m-%d'),
          'end_date': self._end_date.strftime('%Y-%m-%d'),
          'strategy_name': self._strat_info.strategy_name,
          'machine': machine,
          'market_type': container.account_request.market_type,
          'exchange': container.account_request.exchange,
          'owner': container.account_request.owner,
          'symbol': container.symbol
      }
      filename = 'order_plot.html'
      with open('%s/%s' % (output_dir, filename), mode='w') as fd:
        script = bokeh.embed.server_document(server_url, arguments=arguments)
        fd.write(html % script)

  def get_history(self):
    return {key: container.get_history() for key, container in
            self._oe_containers.items()}


def get_order_event_data_dir(
    *, root_dir, start_date, end_date, strategy_name, machine, account_request, symbol):
  config = DailyStrategyPlotterConfig(start_date, end_date, root_dir)
  output_dir = config.get_order_plot_dir(
          strategy_name=strategy_name,
          machine=machine,
          account_request=account_request,
          symbol=symbol)
  output_dir = os.path.join(output_dir, 'data')
  return output_dir


def _gen_order_dict(order_event):
  assert order_event.HasField('event_time')
  assert order_event.HasField('type')
  assert order_event.HasField('order_state')
  order_sign = None
  if order_event.HasField('order_side'):
    order_sign = get_order_sign(order_event.order_side)
    assert order_sign != 0
    order_sign = 'BUY_ORDER' if order_sign > 0 else 'SELL_ORDER'
  order_dict_orig = json_format.MessageToDict(
      order_event, preserving_proto_field_name=True)
  keys = [
      'type', 'order_price', 'order_qty', 'order_side',
      'internal_order_id', 'external_order_id', 'fill_price', 'fill_qty', 'order_state']
  order_dict = {key: order_dict_orig.get(key) for key in keys}
  # proto field with type of int64 is converted to string in MessageToDict()
  order_dict['event_timestamp'] = int(order_event.event_time)
  internal_order_id = order_dict['internal_order_id']
  if order_event.HasField('proc_order_id'):
    assert internal_order_id is None
    internal_order_id = str(order_event.proc_order_id)
  assert internal_order_id is not None
  order_dict['internal_order_id'] = internal_order_id
  order_dict['order_sign'] = order_sign
  return order_dict


def _launch_replay_func(*,
                        strat_info,
                        start_date,
                        end_date,
                        kafka_config,
                        strat_log_root_dir,
                        output_dir,
                        server_url=None):
  logging.info('start replay log')
  strategy_name = strat_info.strategy_name
  machine = strat_info.machine
  start_time = datetime.datetime.combine(
      start_date, datetime.time.min.replace(tzinfo=pytz.UTC))
  end_time = datetime.datetime.combine(
      end_date, datetime.time.max.replace(tzinfo=pytz.UTC))
  replayer = OrderLogReplayer(strat_info, start_date, end_date)
  run_from_strat_log_archive(on_log_callback=replayer.on_log,
                             start_time=start_time,
                             end_time=end_time,
                             root_dir=os.path.expanduser(strat_log_root_dir),
                             machine=machine,
                             strategy_name=strategy_name)
  replayer.dump_history(output_dir)


def _dump_order_event(*,
                      start_date,
                      end_date,
                      strategy_list,
                      kafka_config,
                      strat_log_root_dir,
                      output_dir,
                      server_url,
                      max_workers):
  func_list = []
  strat_infos = get_strat_info_list(
      strat_log_root_dir=strat_log_root_dir,
      kafka_config=kafka_config,
      strategy_list=strategy_list,
      start_date=start_date,
      end_date=end_date)
  for strat_info in strat_infos:
    replay_func = functools.partial(
        _launch_replay_func,
        strat_info=strat_info,
        start_date=start_date,
        end_date=end_date,
        kafka_config=kafka_config,
        strat_log_root_dir=strat_log_root_dir,
        output_dir=output_dir,
        server_url=server_url)
    func_list.append((strat_info, replay_func))
  with ProcessPoolExecutor(max_workers=max_workers) as executor:
    future_list = [(key, executor.submit(func)) for key, func in func_list]
  success = 0
  for key, future in future_list:
    try:
      future.result()
      success += 1
    except Exception as e:
      logging.error('Fail to dump strategy log. %s, %s, %s', key._asdict(), type(e), e)
      logging.error(traceback.format_exc())
  logging.info('Success rate: %s/%s' % (success, len(future_list)))


def main(_):
  start_date = FLAGS.start_date
  assert start_date, '--start_date must be specified.'
  end_date = FLAGS.end_date
  assert end_date, '--end_date must be specified.'
  start_date = convert_string_to_dates(start_date)[0]
  end_date = convert_string_to_dates(end_date)[0]
  trading_dates = [td for td in iterate_date(start_date, end_date)]
  kafka_config_filename = FLAGS.kafka_config_filename
  kafka_config = None
  if kafka_config_filename is not None:
    kafka_config = KafkaConfig.from_cmd_config(kafka_config_filename)
  strategy_config_filename = FLAGS.strategy_config_filename
  strategy_universe = query_strategies(as_proto=True)
  strat_list = get_strat_from_json_config(strategy_config_filename)
  if FLAGS.strategy_group is not None:
    strategy_groups = [elem.strip() for elem in FLAGS.strategy_group.split(',')]
    all_strats = [elem.strategy_name for elem in strategy_universe if
                  elem.strategy_group in strategy_groups]
    strat_list = set(strat_list).intersection(set(all_strats))


  for trading_date in trading_dates:
    logging.info('\nRunning for %s' % trading_date.strftime('%Y%m%d'))
    _dump_order_event(
        start_date=trading_date,
        end_date=trading_date,
        strategy_list=strat_list,
        kafka_config=kafka_config,
        strat_log_root_dir=FLAGS.root_dir,
        output_dir=FLAGS.output_dir,
        server_url=FLAGS.server_url,
        max_workers=FLAGS.max_workers)


if __name__ == '__main__':
  flags.DEFINE_string('start_date', None, 'Start date in form of %Y%m%d.')
  flags.DEFINE_string('end_date', None, 'End date in form of %Y%m%d.')
  flags.DEFINE_string('kafka_config_filename', None, 'kafka config')
  flags.DEFINE_string('root_dir', None, '')
  flags.DEFINE_string('output_dir', None, '')
  flags.DEFINE_string('server_url', None, '')
  flags.DEFINE_integer('max_workers', 24, 'Max number of workers.')
  flags.DEFINE_string('strategy_group', None, 'filter by strategy_group')
  flags.DEFINE_string('strategy_config_filename',
                      '../../coin_deploy/support_monitor/config/strat_config/strat_all.json',
                      'strategy_config_filename')

  app.run(main)
